#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

//# written by wingsummer

#define THREAD_TABLE_MAX_SIZE (5)
#define STACK_SIZE (1024 * 1024)

typedef void (*thread_pointer)();
typedef enum bool { false, true } bool;

typedef enum thread_state { DEAD, RUNNING, SLEEP } thread_state;
typedef struct thread {
    void *stack;
    void *stackpc;
    thread_pointer pc;
    thread_state state;
} thread;

thread thread_table[THREAD_TABLE_MAX_SIZE];
thread *current_thread = &thread_table[0];

bool create_thread(thread_pointer func) {
    if (func == NULL) {
        puts("create_thread failed: please input invalid excution address");
        return false;
    }

    for (int i = 0; i < THREAD_TABLE_MAX_SIZE; i++) {
        if (thread_table[i].state != DEAD) {
            continue;
        }
        thread *th = &thread_table[i];
        th->pc = func;
        th->stack = (void *)((char *)(malloc(STACK_SIZE)) + STACK_SIZE);
        th->stackpc = th->stack;
        if (th->stack == NULL) {
            puts("create_thread failed: malloc thread stack");
            return false;
        }
        th->state = SLEEP;
        return true;
    }
    return false;
}

static int swap_context_i = 0;
#define DECLARE_STRUCT_OFFSET(type, member) [member] "i"(offsetof(type, member))

__attribute__((naked)) __attribute__((fastcall)) void
swap_context(thread *old, thread *new) {
    asm volatile("pushq %rax;"
                 "pushq %rbx;"
                 "pushq %rcx;"
                 "pushq %rdx;"
                 "pushq %rbp;"
                 "pushq %rsi;"
                 "pushq %r8;"
                 "pushq %r9;"
                 "pushq %r10;"
                 "pushq %r11;"
                 "pushq %r12;"
                 "pushq %r13;"
                 "pushq %r14;"
                 "pushq %r15;"
                 "pushfq");

    // rcx: old, rdx: new
    asm volatile("movq %%rsp,%c[stackpc](%%rcx);" ::DECLARE_STRUCT_OFFSET(
                     thread, stackpc)
                 : "rcx", "rdx");
    asm volatile("movq %0,%c[state](%%rcx);"
                 "incq %1;" ::"i"(SLEEP),
                 "m"(swap_context_i), DECLARE_STRUCT_OFFSET(thread, state)
                 : "rcx", "rdx");

    asm volatile("movq %c[pc](%%rdx),%%rax;"
                 "test %%rax,%%rax;"
                 //> if not first started
                 "jz pcnull%=;"
                 // r8 : function handler
                 "movq %%rax,%%r8;"
                 "xor %%eax,%%eax;"
                 "movq %%rax,%c[pc](%%rdx);"
                 "movq %c[stack](%%rdx),%%rbp;"
                 "movq %%rbp,%%rsp;"
                 // fill up stack info
                 "lea idle%=(%%rip), %%rax;"
                 "push %%rax;"
                 "callq *%%r8;"
                 "pcnull%=:;"
                 //> if first started
                 "movq %1,%c[state](%%rdx);"
                 "movq %c[stackpc](%%rdx),%%rsp;"
                 "jmp sw%=;"

                 "idle%=:;"
                 "call %P0;"
                 "jmp idle%=;"
                 "sw%=:;" ::"i"(swap_context),
                 "i"(RUNNING), DECLARE_STRUCT_OFFSET(thread, state),
                 DECLARE_STRUCT_OFFSET(thread, pc),
                 DECLARE_STRUCT_OFFSET(thread, stack),
                 DECLARE_STRUCT_OFFSET(thread, stackpc)
                 : "rax", "rcx", "rdx");

    asm volatile("popfq;"
                 "popq %r15;"
                 "popq %r14;"
                 "popq %r13;"
                 "popq %r12;"
                 "popq %r11;"
                 "popq %r10;"
                 "popq %r9;"
                 "popq %r8;"
                 "popq %rsi;"
                 "popq %rbp;"
                 "popq %rdx;"
                 "popq %rcx;"
                 "popq %rbx;"
                 "popq %rax;"
                 "ret;");
}

void swap_context_caller() {
    thread *new_thread = NULL;
    thread *old_thread = NULL;
    bool notfounded = true;
    while (true) {
        for (; swap_context_i < THREAD_TABLE_MAX_SIZE; swap_context_i++) {
            if (thread_table[swap_context_i].state == SLEEP) {
                old_thread = current_thread;
                new_thread = &thread_table[swap_context_i];
                current_thread = new_thread;
                swap_context(old_thread, new_thread);
                notfounded = false;
                break;
            }
        }
        if (notfounded) {
            swap_context_i = 0;
        } else {
            break;
        }
    }
}

void print(const char *message) {
    puts(message);
    swap_context_caller();
}

//======================================================

#include <signal.h>
static volatile sig_atomic_t keep_running = 1;

#ifdef __WINNT
#include <windows.h>
BOOL WINAPI ctlevent_handler(DWORD dwCtrlType) {
    if (dwCtrlType == CTRL_C_EVENT) {
        keep_running = 0;
        return TRUE;
    }
    return FALSE;
}
#else
static void sig_handler(int _) {
    (void)_;
    keep_running = 0;
}
#endif

//======================================================

void thread_test_1() {
    while (true) {
        print("thread_test_1");
    }
}

void thread_test_2() {
    while (true) {
        print("thread_test_2");
    }
}

void thread_test_3() {
    while (true) {
        print("thread_test_3");
    }
}

int main() {

// Ctrl+C handler
#ifdef __WINNT
    SetConsoleCtrlHandler((PHANDLER_ROUTINE)ctlevent_handler, TRUE);
#else
    signal(SIGINT, sig_handler);
#endif

    // init thread table
    for (int i = 0; i < THREAD_TABLE_MAX_SIZE; i++) {
        memset(&thread_table[i], 0, sizeof(thread));
    }
    thread *th0 = &thread_table[0];
    th0->pc = NULL;
    th0->stackpc = NULL;
    asm("movq %%rbp,%0" ::"m"(th0->stack));
    th0->state = RUNNING;

    // create thread
    create_thread(thread_test_1);
    create_thread(thread_test_2);
    create_thread(thread_test_3);

    // start coroutine test
    while (keep_running) {
        print("main");
    }

    puts("[Interrput] receiving Ctrl+C, cleaning up...");

    // clean up
    for (int i = 1; i < THREAD_TABLE_MAX_SIZE; ++i) {
        void *stack = thread_table[i].stack;
        if (stack) {
            free(stack);
        }
    }

    return 0;
}
